// Author: Ivan Chau
package uk.ac.cam.kcic2.photonMap;

import uk.ac.cam.kcic2.math.Params;
import uk.ac.cam.kcic2.math.RotationY;
import uk.ac.cam.kcic2.math.RotationZ;
import uk.ac.cam.kcic2.math.Vector3D;
import uk.ac.cam.kcic2.rayTracer.RayTracer;
import uk.ac.cam.kcic2.scene.IntersectionInfo;
import uk.ac.cam.kcic2.scene.Primitives;

public class Photon {
	public Vector3D position;
	public final Vector3D power;
	// power packed as 3 doubles using RGB-format
	public final Vector3D direction;
	// compressed incident direction
	private short plane = -1;
	// 2 bits flag used in kd-tree
	public final boolean direct;

	public Photon(double[] pos, double[] dir, double[] pow, boolean d){
		position = new Vector3D(pos);
		direction = new Vector3D(dir);
		power = new Vector3D(pow);
		direct = d;
	}

	public Photon(Vector3D pos, Vector3D dir, Vector3D pow, boolean d){
		position = pos;
		direction = dir;
		power = pow;
		direct = d;
	}

	public double getPosition(int index){return position.get(index);}
	public double getPower(int index){return power.get(index);}
	public short getPlane(){return plane;}
	public void setPlane(short flag){plane = flag;}

	public Photon hitSurface(IntersectionInfo intersection, 
			PhotonMap pMap, boolean caustics){	// Russian Roulette

		Primitives hitObject = intersection.object;
		Vector3D diffuseCoeffs = hitObject.colourCoeffs.diffuseCoeffs;
		Vector3D specularCoeffs = 
				hitObject.colourCoeffs.specularCoeffs;
		Vector3D transmittance = hitObject.colourCoeffs.transmittance;
		double refractiveIndex = 
				hitObject.colourCoeffs.refractiveIndex;
		Vector3D normal = intersection.normal;
		Vector3D hitPoint = intersection.position;

		Vector3D reflectCoeffs = diffuseCoeffs.plus(specularCoeffs);
		double probR = Math.max(reflectCoeffs.x, 
				Math.max(reflectCoeffs.y, reflectCoeffs.z));
		double probDR = 
				(diffuseCoeffs.x + diffuseCoeffs.y + diffuseCoeffs.z)/
				(reflectCoeffs.x + reflectCoeffs.y + reflectCoeffs.z)*
				probR;
		double probSR = (specularCoeffs.x + specularCoeffs.y + 
				specularCoeffs.z)/
				(reflectCoeffs.x + reflectCoeffs.y + reflectCoeffs.z)*
				probR;
		assert Math.abs(probSR + probDR - probR) < Params.EPILSON;
		double probT = Math.max(transmittance.x, 
				Math.max(transmittance.y, transmittance.z));

		assert(probDR + probSR + probT <= 1.0);

		double randomVar = Math.random();
		if (randomVar < probDR){	// photon reflects diffusely
			if (caustics && direct){return null;}

			if(!direct){
				Photon q = new Photon(hitPoint, this.direction, 
						this.power, direct);
				pMap.insert(q);
			}
			if (caustics){
				return null;
			}

			// pick a random direction in the hemisphere above the 
			// intersection point
			Vector3D newDirection;

			// rotate the axis to the xz plane
			double theta = Math.acos(normal.z);
			double phi = Math.atan2(normal.y, normal.x);
			RotationZ r1 = new RotationZ(-phi);
			newDirection = r1.mult(normal);
			// rotation the axis to align with the z-axis
			RotationY r2 = new RotationY(-theta);
			newDirection = r2.mult(newDirection);
			assert newDirection.sub(new Vector3D(0,0,1)).norm() < 
			Params.EPILSON;
			// rotate the normal to the new direction
			theta = Math.acos(Math.sqrt(Math.random()));
			phi = Math.PI*2*Math.random();
			RotationY r3 = new RotationY(theta);
			newDirection = r3.mult(newDirection);
			RotationZ r4 = new RotationZ(phi);
			newDirection = r4.mult(newDirection);
			// reverse the appropriate transform back to the hitPoint
			newDirection = r2.inverse().mult(newDirection);
			newDirection = r1.inverse().mult(newDirection);

			// scale the spectral power of the photon
			double red = this.power.x*diffuseCoeffs.x/probDR;
			double green = this.power.y*diffuseCoeffs.y/probDR;
			double blue = this.power.z*diffuseCoeffs.z/probDR;

			newDirection = newDirection.normalize();
			newDirection = newDirection.mult(-1);

			Photon p = new Photon(hitPoint, newDirection, 
					new Vector3D(red, green, blue), false);

			return p;

		}
		else if (randomVar < probDR + probSR){	
			// photon reflects specularly
			if (caustics){return null;}
			Vector3D reflectionDirection = normal.mult(
					2*this.direction.dot(normal)).sub(this.direction);

			// scale the spectral power of the photon
			double red = this.power.x*specularCoeffs.x/probSR;
			double green = this.power.y*specularCoeffs.y/probSR;
			double blue = this.power.z*specularCoeffs.z/probSR;

			reflectionDirection = reflectionDirection.mult(-1);

			Photon p = new Photon(hitPoint, reflectionDirection, 
					new Vector3D(red, green, blue), false);

			return p;
		}
		else if (randomVar < probDR + probSR + probT){	
			// photon refracts
			if (!caustics && direct){return null;}
			double refractiveIndicesRatio = 1/refractiveIndex;

			if (intersection.inwards == false){
				refractiveIndicesRatio = refractiveIndex;
			}
			Vector3D refractedDirection = RayTracer.refract(
					this.direction.mult(-1), normal, 
					refractiveIndicesRatio);

			// scale the spectral power of the photon
			double red = this.power.x*transmittance.x/probT;
			double green = this.power.y*transmittance.y/probT;
			double blue = this.power.z*transmittance.z/probT;

			refractedDirection = refractedDirection.mult(-1);

			Photon p = new Photon(hitPoint, refractedDirection, 
					new Vector3D(red, green, blue), false);

			return p;
		}
		else{	// photon is absorbed
			return null;
		}
	}
}
